#!/usr/bin/python

import sys, os

# From the squeue man page
states = {
    'PD' : 'pending',
    'R' : 'running', 
    'CA' : 'cancelled',
    'CF' : 'configuring',
    'CG' : 'completing',
    'CD' : 'completed',
    'F' : 'failed',
    'TO' : 'timeout',
    'NF' : 'node failure'
}

if len(sys.argv) != 5:
    print "check_slurm_custom -w <warning> -c <critical>"
    sys.exit(3)

W = int(sys.argv[2])
C = int(sys.argv[4])

cmd = "/usr/local/bin/squeue -h --format='%u' | sort | uniq"
users = [u.strip() for u in os.popen(cmd).readlines()]

cmd = "/usr/local/bin/squeue -h --format='%t %u'"
lines = [l.strip() for l in os.popen(cmd).readlines()]
counts = {}
for line in lines: 
    state, user = line.split()
    state = states.get(state, 'unknown')
    counts[user] = counts.get(user, {})
    counts[user][state] = counts[user].get(state, 0) + 1

total = sum([sum(v.values()) for k, v in counts.iteritems()])
entries = [
    ",".join( [
        str(ele) for ele in [
            k, v.get('running', 0), v.get('pending', 0), v.get('unknown', 0)
        ]
    ] ) for k, v in counts.iteritems()
]

output = ":".join(entries)

if total < W:
    print "SLURM_CUSTOM OK - (%i jobs)|%s" % (total, output)
    sys.exit(0)
if total < C:
    print "SLURM_CUSTOM WARN - (%i jobs)|%s" % (total, output)
    sys.exit(1)

print "SLURM_CUSTOM CRIT - (%i jobs)|%s" % (total, output)
sys.exit(2)

